import argparse
import math
import os

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as transforms
from sklearn.metrics import accuracy_score
from torch.autograd import Variable
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.utils import save_image
from vit_pytorch import ViT

from AutoGPU import autoGPU

autoGPU(1, 'auto')
os.makedirs("images", exist_ok=True)

parser = argparse.ArgumentParser()
parser.add_argument("--n_epochs", type=int, default=200, help="number of epochs of training")
parser.add_argument("--batch_size", type=int, default=256, help="size of the batches")
parser.add_argument("--lr", type=float, default=0.0002, help="adam: learning rate")
parser.add_argument("--b1", type=float, default=0.5, help="adam: decay of first order momentum of gradient")
parser.add_argument("--b2", type=float, default=0.999, help="adam: decay of first order momentum of gradient")
parser.add_argument("--n_cpu", type=int, default=8, help="number of cpu threads to use during batch generation")
parser.add_argument("--latent_dim", type=int, default=100, help="dimensionality of the latent space")
parser.add_argument("--img_size", type=int, default=28, help="size of each image dimension")
parser.add_argument("--channels", type=int, default=1, help="number of image channels")
parser.add_argument("--sample_interval", type=int, default=50, help="interval betwen image samples")
opt = parser.parse_args()
print(opt)

img_shape = (opt.channels, opt.img_size, opt.img_size)

cuda = True if torch.cuda.is_available() else False

# from models.Celeba64_TransGAN import Discriminator
from myTrans2 import Discriminator, Generator

# Loss function
adversarial_loss = torch.nn.BCELoss()
classification_loss = torch.nn.CrossEntropyLoss()

# Initialize generator and discriminator

generator = Generator()
discriminator = Discriminator()

print("resnet50 have {} paramerters in total".format(sum(x.numel() for x in generator.parameters())))

if cuda:
    generator.cuda()
    discriminator.cuda()
    adversarial_loss.cuda()
    classification_loss.cuda()


v = ViT(
    image_size = 256,
    patch_size = 32,
    num_classes = 1000,
    dim = 1024,
    depth = 6,
    heads = 16,
    mlp_dim = 2048,
    dropout = 0.1,
    emb_dropout = 0.1
)


# Configure data loader
os.makedirs("../../data/mnist", exist_ok=True)
dataloader = torch.utils.data.DataLoader(
    datasets.MNIST(
        "../../data/mnist",
        train=True,
        download=True,
        transform=transforms.Compose(
            [transforms.Resize(opt.img_size), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])]
        ),
    ),
    batch_size=opt.batch_size,
    shuffle=True,
)

# Optimizers
optimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))
optimizer_D = torch.optim.Adam(discriminator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))

Tensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor

# ----------
#  Training
# ----------




for epoch in range(opt.n_epochs):
    trainloss_sigma = 0.0 
    for i, (imgs, gt) in enumerate(dataloader):
        one_hot_labels = torch.FloatTensor(gt.shape[0], 10).cuda()
        # Adversarial ground truths6
        valid = Variable(Tensor(imgs.size(0), 1).fill_(1.0), requires_grad=False)
        fake = Variable(Tensor(imgs.size(0), 1).fill_(0.0), requires_grad=False)

        # Configure input
        real_imgs = Variable(imgs.type(Tensor))

        # -----------------
        
        # -----------------

        optimizer_G.zero_grad()

        # Sample noise as generator input
        z = torch.cuda.FloatTensor(np.random.normal(0, 1, (imgs.shape[0], 90))).cuda()
        one_hot_labels.zero_()
        one_hot_labels.scatter_(1, gt.to('cuda').view(gt.shape[0],1), 1)


        # Generate a batch of images
        gen_imgs = generator(z, one_hot_labels)
        
        # Loss measures generator's ability to fool the discriminator
        valid_gen, gt_gen = discriminator(gen_imgs)
        g_loss = adversarial_loss(valid_gen, valid)
        # c_loss = classification_loss(gt_gen, torch.argmax(one_hot_labels, -1))
        g_loss.backward(retain_graph=True)
        # c_loss.backward()
        optimizer_G.step()
     
        # # ---------------------
        # #  Train Discriminator
        # # ---------------------

        optimizer_D.zero_grad()

        # # Measure discriminator's ability to classify real from generated samples
        reality_real, cls_real = discriminator(real_imgs)
        reality_real_loss = adversarial_loss(reality_real, valid)
        cls_real_loss = classification_loss(cls_real, torch.argmax(one_hot_labels, -1))
        gen_imgs1 = torch.rand(gen_imgs.shape).cuda()
        with torch.no_grad():
            gen_imgs1.set_(gen_imgs)

        reality_fake, cls_fake = discriminator(gen_imgs1)
        reality_fake_loss = adversarial_loss(reality_fake, fake)
        cls_fake_loss = classification_loss(cls_fake, torch.argmax(one_hot_labels, -1))

        d_loss = (reality_real_loss+ reality_fake_loss) / 2

        d_loss.backward(retain_graph=True)
        cls_fake_loss.backward(retain_graph=True)
        cls_real_loss.backward()
        optimizer_D.step()

    loss_avg = trainloss_sigma / len(dataloader)
    with torch.no_grad():
        y_pre = []
        y_true = []
        for batch_idx, data in enumerate(dataloader):
            for idx, item in enumerate(data):
                data[idx] = item.to('cuda')
            *images, labels = data
            discriminator.eval()
            _, outputs = discriminator(*images)  # forward
            #outputs.detach_()  # 不求梯度
            _, predicted = torch.max(outputs.data, 1)  # 统计
            # 统计混淆矩阵
            y_pre += list(predicted.cpu().numpy())
            y_true += list(labels.cpu().numpy())
    print(accuracy_score(y_true, y_pre)
)
        # print(
        #     "[Epoch %d/%d] [Batch %d/%d] [G loss: %f]"
        #     % (epoch, opt.n_epochs, i, len(dataloader),  g_loss.item())
        # )

        # print(
        #     "[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f]"
        #     % (epoch, opt.n_epochs, i, len(dataloader), d_loss.item(), g_loss.item())
        # )

        # batches_done = epoch * len(dataloader) + i
        # if batches_done % opt.sample_interval == 0:
        #     save_image(gen_imgs.data[:25], "images/%d.png" % batches_done, nrow=5, normalize=True)


    
